from PIL import Image, ImageGrab
from gravityassist.api import client, retry_api
from gravityassist.image import encode_image_to_url
from gravityassist.config import get_options

def describe_image(path: str | Image.Image, prompt: str = "What’s in this image?", max_tokens: int = 512, max_size: int = 512):
    if isinstance(path, Image.Image):
        image_url = encode_image_to_url(image, max_size)
    elif path.startswith('https://') or path.startswith('http://'):
        image_url = path
    else:
        image_url = encode_image_to_url(Image.open(path), max_size)
    response = retry_api(client.chat.completions.create,
        model="gpt-4-vision-preview",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": prompt,
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_url,
                        },
                    },
                ],
            },
        ],
        temperature=get_options()["temperature"],
        max_tokens=max_tokens,
    )
    description = response.choices[0].message.content
    print(description)
    return {'description': description}

if __name__ == '__main__':
    from gravityassist.config import load_persistent_data
    load_persistent_data('instance', False)
    image = ImageGrab.grabclipboard()
    # image.show()
    # print(describe_image(image, '这幅图片里有什么？'))
    print(describe_image(image, '基于图片写一个美丽的故事', 256))
